#include <pcap.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>

#define SNAPLEN 65535
#define PROMISC 1
#define TIMEOUT 500

pcap_t* init_pcap(const char* dev) {
    char errbuf[PCAP_ERRBUF_SIZE];
    pcap_t* handle = pcap_open_live(dev, SNAPLEN, PROMISC, TIMEOUT, errbuf);
    
    struct bpf_program fp;
    char filter_exp[] = "ip";  // 只捕获IP流量
    pcap_compile(handle, &fp, filter_exp, 0, PCAP_NETMASK_UNKNOWN);
    pcap_setfilter(handle, &fp);
    
    return handle;
}

typedef struct {
    uint32_t src_ip;
    uint32_t dst_ip;
    uint16_t src_port;
    uint16_t dst_port;
    uint8_t protocol;  // IPPROTO_TCP/UD等
    
    // 流量统计
    uint64_t total_rx;
    uint64_t total_tx;
    uint64_t peak_rx;
    uint64_t peak_tx;
    
    // 滑动窗口统计
    struct {
        uint64_t bytes[40];
        int current;
    } window;
} Connection;

// 使用哈希表加速查找
#define HASH_SIZE 4096
Connection* connection_table[HASH_SIZE];

void packet_handler(u_char *user, const struct pcap_pkthdr *h, const u_char *bytes) {
    struct ip *ip_header = (struct ip*)(bytes + 14);  // 跳过以太网头
    struct tcphdr *tcp_header = (struct tcphdr*)(bytes + 14 + (ip_header->ip_hl<<2));

    // 生成连接标识
    uint32_t src = ip_header->ip_src.s_addr;
    uint32_t dst = ip_header->ip_dst.s_addr;
    uint16_t sport = ntohs(tcp_header->th_sport);
    uint16_t dport = ntohs(tcp_header->th_dport);

    // 计算哈希键
    int key = (src + dst + sport + dport) % HASH_SIZE;
    
    // 查找或创建连接记录
    Connection *conn = find_or_create_connection(key, src, dst, sport, dport);
    
    // 更新统计
    uint64_t pkt_len = h->len - 14;  // 排除以太网头
    update_stats(conn, pkt_len, src == local_ip);
}

Connection* find_or_create_connection(int key, uint32_t src, uint32_t dst, 
                                     uint16_t sport, uint16_t dport) {
    Connection *curr = connection_table[key];
    while(curr) {
        if(curr->src_ip == src && curr->dst_ip == dst &&
           curr->src_port == sport && curr->dst_port == dport)
            return curr;
        curr = curr->next;
    }
    
    // 创建新连接
    Connection *new_conn = calloc(1, sizeof(Connection));
    /* 初始化字段... */
    new_conn->next = connection_table[key];
    connection_table[key] = new_conn;
    return new_conn;
}

void update_stats(Connection *conn, uint64_t size, int is_inbound) {
    // 累计总量
    if(is_inbound) {
        conn->total_rx += size;
        conn->window.bytes[conn->window.current] += size;
    } else {
        conn->total_tx += size;
        conn->window.bytes[conn->window.current] += size;
    }

    // 更新峰值
    if(size > (is_inbound ? conn->peak_rx : conn->peak_tx)) {
        if(is_inbound) 
            conn->peak_rx = size;
        else
            conn->peak_tx = size;
    }
}

// 每秒触发窗口滑动
void rotate_window(int sig) {
    for(int i=0; i<HASH_SIZE; i++) {
        Connection *curr = connection_table[i];
        while(curr) {
            curr->window.current = (curr->window.current + 1) % 40;
            curr->window.bytes[curr->window.current] = 0;
            curr = curr->next;
        }
    }
    alarm(1);  // 重新设置定时器
}

uint64_t calculate_average(Connection *conn, int seconds, int is_rx) {
    uint64_t sum = 0;
    int window_pos = conn->window.current;
    
    for(int i=0; i<seconds; i++) {
        int idx = (window_pos - i + 40) % 40;
        sum += conn->window.bytes[idx];
    }
    return sum / seconds;
}

int main() {
    pcap_t *handle = init_pcap("br-lan");  // OpenWrt默认桥接接口
    signal(SIGALRM, rotate_window);
    alarm(1);

    // 开始抓包循环
    pcap_loop(handle, -1, packet_handler, NULL);
    
    pcap_close(handle);
    return 0;
}